import cv2
import numpy as np
import random
import heapq

# 蛇行 但是不考虑下一个目标点

# 游戏参数
width, height = 300, 300  # 地图的宽度和高度
snake_size = 10  # 蛇身和食物的大小
fps = 60  # 游戏帧率，控制蛇的移动速度

# 初始化游戏状态
snake = [(100, 100), (90, 100), (80, 100)]  # 初始蛇身位置，每个元素是一个(x, y)坐标对
direction = 'RIGHT'  # 初始方向
score = 0  # 得分
food = (random.randint(0, width // snake_size - 1) * snake_size,  # 随机生成食物的位置
        random.randint(0, height // snake_size - 1) * snake_size)

def draw_snake_and_food(frame):
    """
    在给定的帧上绘制蛇和食物。
    :param frame: 当前游戏画面的numpy数组
    """
    # 绘制蛇的身体，头部用红色，身体其他部分用绿色
    for i, (x, y) in enumerate(snake):
        color = (0, 255, 0) if i > 0 else (255, 0, 0)  # 蛇的身体用绿色，头部用红色
        cv2.rectangle(frame, (x, y), (x + snake_size, y + snake_size), color, -1)
    
    # 绘制食物，使用红色方块
    cv2.rectangle(frame, (food[0], food[1]), (food[0] + snake_size, food[1] + snake_size), (0, 0, 255), -1)

def move_snake(new_head):
    """
    移动蛇的位置。
    :param new_head: 新的蛇头位置
    """
    global food, score  # 允许函数内部修改全局变量food和score
    if new_head == food:  # 如果新头部位置等于食物位置，则吃掉食物并生成新的食物
        snake.insert(0, new_head)  # 增加蛇的长度
        score += 1  # 增加得分
        food = (random.randint(0, width // snake_size - 1) * snake_size,  # 重新生成食物
                random.randint(0, height // snake_size - 1) * snake_size)
    else:
        snake.insert(0, new_head)  # 移动蛇头
        snake.pop()  # 移除蛇尾

def check_collision(new_head):
    """
    检查蛇是否撞到墙壁或自身。
    :param new_head: 新的蛇头位置
    :return: 如果发生碰撞，返回True；否则返回False
    """
    x, y = new_head
    if x < 0 or x >= width or y < 0 or y >= height:  # 检查是否撞墙
        print(f"Game over: Snake hit the wall at position {new_head}")
        return True
    if new_head in snake:  # 检查是否撞到自身
        print(f"Game over: Snake collided with itself at position {new_head}. Current snake positions: {snake}")
        return True
    return False

def heuristic(a, b):
    """
    A*算法中的启发式函数，计算两点之间的曼哈顿距离。
    :param a: 第一个点
    :param b: 第二个点
    :return: 曼哈顿距离
    """
    return abs(a[0] - b[0]) + abs(a[1] - b[1])

def a_star_search(start, target, obstacles):
    """
    使用A*算法搜索从起始点到目标点的最短路径。
    :param start: 起始点
    :param target: 目标点
    :param obstacles: 障碍物列表
    :return: 从起始点到目标点的路径，如果找不到路径则返回None
    """
    open_set = []
    heapq.heappush(open_set, (0, start, []))  # 将起始点加入开放集
    g_score = {start: 0}  # 从起点到当前节点的实际成本
    f_score = {start: heuristic(start, target)}  # 从起点到当前节点再到目标节点的估计总成本
    came_from = {}  # 记录路径

    while open_set:
        _, current, path = heapq.heappop(open_set)  # 取出成本最低的节点

        if current == target:  # 如果到达目标，则返回路径
            return path + [current]

        for dx, dy in [(-snake_size, 0), (snake_size, 0), (0, -snake_size), (0, snake_size)]:
            neighbor = (current[0] + dx, current[1] + dy)
            tentative_g_score = g_score[current] + 1

            # 检查邻居节点是否有效
            if 0 <= neighbor[0] < width and 0 <= neighbor[1] < height and neighbor not in obstacles:
                if neighbor not in g_score or tentative_g_score < g_score[neighbor]:
                    came_from[neighbor] = current
                    g_score[neighbor] = tentative_g_score
                    f_score[neighbor] = tentative_g_score + heuristic(neighbor, target)
                    heapq.heappush(open_set, (f_score[neighbor], neighbor, path + [current]))

    return None  # 如果没有找到路径，返回None

def get_surrounding_free_spaces(snake):
    """
    获取蛇头周围的空闲格子。
    :param snake: 蛇的位置列表
    :return: 周围的空闲格子列表
    """
    snake_set = set(snake)
    free_spaces = []
    for dx, dy in [(-snake_size, 0), (snake_size, 0), (0, -snake_size), (0, snake_size)]:
        neighbor = (snake[0][0] + dx, snake[0][1] + dy)
        if 0 <= neighbor[0] < width and 0 <= neighbor[1] < height and neighbor not in snake_set:
            free_spaces.append(neighbor)
    return free_spaces

def choose_best_free_space(free_spaces):
    """
    从空闲格子中选择最优的一个。
    :param free_spaces: 空闲格子列表
    :return: 最优的空闲格子
    """
    if not free_spaces:
        return None
    # 选择离食物最近的空闲格子
    best_space = min(free_spaces, key=lambda space: heuristic(space, food))
    return best_space

def draw_path(frame, path):
    """
    在给定的帧上绘制路径。
    :param frame: 当前游戏画面的numpy数组
    :param path: 路径列表
    """
    if path:
        for point in path:
            # 绘制灰色实线框，大小与蛇头一致
            cv2.rectangle(frame, (point[0], point[1]), (point[0] + snake_size, point[1] + snake_size), (128, 128, 128), 1)

def show_message(message, frame=None, font_scale=1, thickness=2, wait_key=True):
    """
    在给定的帧上显示消息，并等待按键。
    :param message: 要显示的消息
    :param frame: 背景帧，如果为None，则创建一个黑色背景
    :param font_scale: 字体大小
    :param thickness: 字体厚度
    :param wait_key: 是否等待按键
    """
    if frame is None:
        frame = np.zeros((height, width, 3), dtype=np.uint8)
    text_size = cv2.getTextSize(message, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)[0]
    text_x = (width - text_size[0]) // 2
    text_y = (height + text_size[1]) // 2
    cv2.putText(frame, message, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thickness)
    cv2.imshow('Snake Game', frame)
    if wait_key:
        cv2.waitKey(0)

def main():
    global direction, last_frame
    game_over = False
    last_frame = None  # 用于保存最后一帧的画面

    # 显示“按任意键开始”
    show_message("Press any key to start the game...", font_scale=0.5, thickness=2)

    while not game_over:
        # 创建空白画布
        frame = np.zeros((height, width, 3), dtype=np.uint8)

        # 绘制蛇和食物
        draw_snake_and_food(frame)

        # 计算新的头部位置
        head_x, head_y = snake[0]
        if direction == 'UP':
            new_head = (head_x, head_y - snake_size)
        elif direction == 'DOWN':
            new_head = (head_x, head_y + snake_size)
        elif direction == 'LEFT':
            new_head = (head_x - snake_size, head_y)
        elif direction == 'RIGHT':
            new_head = (head_x + snake_size, head_y)

        # 寻找路径
        path = a_star_search(snake[0], food, snake)

        if path is None or len(path) < 2:
            # 如果没有找到路径，选择一个最优的空闲格子
            free_spaces = get_surrounding_free_spaces(snake)
            best_space = choose_best_free_space(free_spaces)
            if best_space:
                new_head = best_space
            else:
                game_over = True
                last_frame = frame.copy()
                print("Game over: No free spaces to move to.")
                break
        else:
            # 计算新头部位置
            new_head = path[1]

        # 检查碰撞
        if check_collision(new_head):
            game_over = True
            last_frame = frame.copy()
            break

        # 移动蛇
        move_snake(new_head)

        # 绘制路径
        draw_path(frame, path)

        # 显示得分
        cv2.putText(frame, f"Score: {score}", (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

        # 显示帧
        cv2.imshow('Snake Game', frame)

        # 处理键盘输入
        key = cv2.waitKey(int(1000 / fps)) & 0xFF
        if key == ord('w') and direction != 'DOWN':
            direction = 'UP'
        elif key == ord('s') and direction != 'UP':
            direction = 'DOWN'
        elif key == ord('a') and direction != 'RIGHT':
            direction = 'LEFT'
        elif key == ord('d') and direction != 'LEFT':
            direction = 'RIGHT'
        elif key == ord('q'):
            break

    # 保存最后一帧的画面
    if last_frame is not None:
        cv2.imwrite('last_frame.png', last_frame)
        print("Last frame saved as last_frame.png")

    # 显示最后一帧的画面
    if last_frame is not None:
        show_message(f"Game Over! Final Score: {score}. Press any key to exit...", frame=last_frame, font_scale=0.5, thickness=2, wait_key=True)

    # 等待按键释放所有窗口
    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()